from typing import Optional

import torch
import numpy as np
from scipy.optimize import minimize
from .hand import HandModel, NUM_JOINTS_PER_HAND, NUM_LANDMARKS_PER_HAND
from .hand_skinning import skin_landmarks, _hand_skinning_transform, _joint_local_transform
from pytorch3d.transforms import so3_exp_map, so3_log_map


def inverse_kinematics(
        hand_model: HandModel,
        target_landmarks: torch.Tensor,
        initial_joint_angles: Optional[torch.Tensor] = None,
        initial_wrist_transform: Optional[torch.Tensor] = None,
        max_iterations: int = 100,
        tolerance: float = 1e-4
) -> tuple:
    """
    逆向运动学：根据landmarks和hand_model，计算出腕关节变换矩阵和其他关节角度

    Args:
        hand_model: 手部模型
        target_landmarks: 目标landmarks位置 [21, 3]
        initial_joint_angles: 初始关节角度 [20]，如果为None则使用零初始化
        initial_wrist_transform: 初始腕关节变换矩阵 [4, 4]，如果为None则使用单位矩阵
        max_iterations: 最大迭代次数
        tolerance: 收敛容忍度

    Returns:
        tuple: (wrist_transform, joint_angles, loss_history)
    """
    device = target_landmarks.device
    dtype = target_landmarks.dtype

    # 初始化参数
    if initial_joint_angles is None:
        joint_angles = torch.zeros(NUM_JOINTS_PER_HAND, device=device, dtype=dtype, requires_grad=True)
    else:
        joint_angles = initial_joint_angles.clone().detach().requires_grad_(True)

    if initial_wrist_transform is None:
        wrist_transform = torch.eye(4, device=device, dtype=dtype)
        wrist_transform = wrist_transform.unsqueeze(0)  # 添加batch维度
    else:
        wrist_transform = initial_wrist_transform.clone().detach()

    # 将wrist_transform转换为可优化的参数
    # 使用6自由度表示（3个平移，3个旋转）
    wrist_translation = wrist_transform[0, :3, 3].clone().detach().requires_grad_(False)
    wrist_rotation = so3_log_map(wrist_transform[0, :3, :3].unsqueeze(0))[0].clone().detach().requires_grad_(True)

    # 优化器
    parameters = [{'params': [joint_angles], 'lr': 0.01},
                  {'params': [wrist_rotation], 'lr': 0.05}]
    optimizer = torch.optim.Adam(parameters)

    loss_history = []

    for iteration in range(max_iterations):
        optimizer.zero_grad()

        # 构建手腕变换矩阵
        rot_mat = so3_exp_map(wrist_rotation.unsqueeze(0))[0]
        current_wrist_transform = torch.eye(4, device=device, dtype=dtype)
        current_wrist_transform[:3, :3] = rot_mat
        current_wrist_transform[:3, 3] = wrist_translation
        current_wrist_transform = current_wrist_transform.unsqueeze(0)  # 添加batch维度

        # 前向运动学计算当前landmarks
        current_landmarks = skin_landmarks(hand_model, joint_angles, current_wrist_transform)
        current_landmarks = current_landmarks.squeeze(0)  # 移除batch维度

        # 计算损失
        loss = torch.mean(torch.sum((current_landmarks - target_landmarks) ** 2, dim=-1))

        # 反向传播
        loss.backward()
        optimizer.step()

        loss_history.append(loss.item())

        # 检查收敛
        if loss.item() < tolerance:
            print(f"收敛于迭代 {iteration}, 损失: {loss.item()}")
            break

        if iteration % 10 == 0:
            print(f"迭代 {iteration}, 损失: {loss.item()}")

    # 构建最终的手腕变换矩阵
    rot_mat = so3_exp_map(wrist_rotation.unsqueeze(0))[0]
    final_wrist_transform = torch.eye(4, device=device, dtype=dtype)
    final_wrist_transform[:3, :3] = rot_mat
    final_wrist_transform[:3, 3] = wrist_translation

    return final_wrist_transform, joint_angles.detach(), loss_history


def inverse_kinematics_analytic(
        hand_model: HandModel,
        target_landmarks: torch.Tensor,
        wrist_position: Optional[torch.Tensor] = None
) -> tuple:
    """
    使用解析方法进行逆向运动学（简化版）
    主要针对手指关节，手腕位置可以通过其他方法估计

    Args:
        hand_model: 手部模型
        target_landmarks: 目标landmarks位置 [21, 3]
        wrist_position: 手腕位置估计 [3,]，如果为None则使用目标landmarks中的手腕位置

    Returns:
        tuple: (wrist_transform, joint_angles)
    """
    device = target_landmarks.device
    dtype = target_landmarks.dtype

    # 获取手腕位置
    if wrist_position is None:
        wrist_position = target_landmarks[5]  # 假设第5个landmark是手腕

    # 初始化关节角度
    joint_angles = torch.zeros(NUM_JOINTS_PER_HAND, device=device, dtype=dtype)

    # 对于每个手指，计算逆向运动学
    for finger_idx in range(5):  # 5个手指
        # 获取手指相关的landmarks索引
        if finger_idx == 0:  # 拇指
            base_idx = 1  # 拇指基部
            joint1_idx = 2  # 拇指第一关节
            joint2_idx = 3  # 拇指第二关节
            tip_idx = 4  # 拇指指尖
        else:  # 其他手指
            base_idx = 5 + (finger_idx - 1) * 4
            joint1_idx = base_idx + 1
            joint2_idx = base_idx + 2
            joint3_idx = base_idx + 3
            tip_idx = base_idx + 4

        # 获取目标位置
        base_pos = target_landmarks[base_idx]
        tip_pos = target_landmarks[tip_idx]

        # 计算手指方向向量
        finger_dir = tip_pos - base_pos
        finger_length = torch.norm(finger_dir)
        finger_dir = finger_dir / finger_length

        # 简化逆向运动学：根据方向向量计算关节角度
        # 这里使用简化方法，实际应用中需要更复杂的计算
        if finger_idx == 0:  # 拇指
            # 拇指有2个关节
            angle1 = torch.atan2(finger_dir[1], finger_dir[0])
            angle2 = torch.atan2(finger_dir[2], torch.sqrt(finger_dir[0] ** 2 + finger_dir[1] ** 2))

            joint_angles[0] = angle1  # 拇指第一关节
            joint_angles[1] = angle2  # 拇指第二关节
        else:
            # 其他手指有3个关节
            angle1 = torch.atan2(finger_dir[1], finger_dir[0])
            angle2 = torch.atan2(finger_dir[2], torch.sqrt(finger_dir[0] ** 2 + finger_dir[1] ** 2))
            angle3 = angle2 * 0.7  # 简化假设：第三关节角度是第二关节的70%

            joint_index = 2 + (finger_idx - 1) * 3
            joint_angles[joint_index] = angle1  # 第一关节
            joint_angles[joint_index + 1] = angle2  # 第二关节
            joint_angles[joint_index + 2] = angle3  # 第三关节

    # 构建手腕变换矩阵
    wrist_transform = torch.eye(4, device=device, dtype=dtype)
    wrist_transform[:3, 3] = wrist_position

    return wrist_transform, joint_angles


def optimize_hand_pose(
        hand_model: HandModel,
        target_landmarks: torch.Tensor,
        initial_pose: Optional[tuple] = None
) -> tuple:
    """
    使用数值优化方法优化手部姿态

    Args:
        hand_model: 手部模型
        target_landmarks: 目标landmarks位置 [21, 3]
        initial_pose: 初始姿态 (wrist_transform, joint_angles)

    Returns:
        tuple: 优化后的 (wrist_transform, joint_angles)
    """
    device = target_landmarks.device
    dtype = target_landmarks.dtype

    # 初始化参数
    if initial_pose is None:
        wrist_transform = torch.eye(4, device=device, dtype=dtype)
        joint_angles = torch.zeros(NUM_JOINTS_PER_HAND, device=device, dtype=dtype)
    else:
        wrist_transform, joint_angles = initial_pose

    # 将参数转换为numpy数组用于scipy优化
    def pack_parameters(wrist_transform, joint_angles):
        # 将手腕变换矩阵转换为6自由度表示（3平移 + 3旋转）
        translation = wrist_transform[:3, 3].detach().cpu().numpy()
        rotation = so3_log_map(wrist_transform[:3, :3].unsqueeze(0))[0].detach().cpu().numpy()
        angles = joint_angles.detach().cpu().numpy()
        return np.concatenate([translation, rotation, angles])

    def unpack_parameters(params):
        translation = torch.tensor(params[:3], device=device, dtype=dtype)
        rotation = torch.tensor(params[3:6], device=device, dtype=dtype)
        angles = torch.tensor(params[6:], device=device, dtype=dtype)

        # 重建手腕变换矩阵
        rot_mat = so3_exp_map(rotation.unsqueeze(0))[0]
        wrist_transform = torch.eye(4, device=device, dtype=dtype)
        wrist_transform[:3, :3] = rot_mat
        wrist_transform[:3, 3] = translation

        return wrist_transform, angles

    # 目标函数
    def objective_function(params):
        wrist_transform, joint_angles = unpack_parameters(params)

        # 前向运动学计算当前landmarks
        current_landmarks = skin_landmarks(
            hand_model,
            joint_angles.unsqueeze(0),
            wrist_transform.unsqueeze(0)
        ).squeeze(0)

        # 计算损失
        loss = torch.mean(torch.sum((current_landmarks - target_landmarks) ** 2, dim=-1))
        return loss.item()

    # 初始参数
    initial_params = pack_parameters(wrist_transform, joint_angles)

    # 优化
    result = minimize(
        objective_function,
        initial_params,
        method='L-BFGS-B',
        options={'maxiter': 100, 'disp': True}
    )

    # 提取优化后的参数
    optimized_wrist_transform, optimized_joint_angles = unpack_parameters(result.x)

    return optimized_wrist_transform, optimized_joint_angles
